package ynu.edu.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import ynu.edu.domain.PageDTO;
import ynu.edu.domain.dto.Text2ImgParameters;
import ynu.edu.domain.dto.Txt2ImgRequest;
import ynu.edu.domain.vo.StableDiffusionApiResponse;
import ynu.edu.domain.vo.Txt2ImgResponse;
import ynu.edu.entity.Text2Img;
import ynu.edu.mapper.Text2ImgMapper;
import ynu.edu.service.Text2ImgService;
import ynu.edu.utils.AliOssUtil;
import ynu.edu.utils.Text2imgConstructionUtil;

import java.time.LocalDateTime;
import java.util.*;

/**
 * Stable Diffusion服务实现类
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class Text2ImgServiceImpl implements Text2ImgService {

    @Value("${stable.diffusion.api.url:http://127.0.0.1:7860}")
    private String apiBaseUrl;

    private final Text2imgConstructionUtil text2imgConstructionUtil;

    private final AliOssUtil aliOssUtil;

    private final RestTemplate restTemplate;

    private final ObjectMapper objectMapper;

    private final Text2ImgMapper text2ImgMapper;

    @Override
    public Txt2ImgResponse txt2img(Txt2ImgRequest request) {
        long startTime = System.currentTimeMillis();
        
        try {
            log.info("开始图片生成，请求: {}", request);
            
            // 1. 检查Stable Diffusion服务健康状态
            if (!checkServiceHealth()) {
                log.error("Stable Diffusion服务不可用");
                return buildErrorResponse("Stable Diffusion服务不可用，请检查服务状态", startTime);
            }
            
            // 2. 构造参数
            Text2ImgParameters params = text2imgConstructionUtil.constructParameters(request);
            log.info("参数构造完成，种子: {}", params.getSeed());
            
            // 3. 发送请求到Stable Diffusion API
            StableDiffusionApiResponse sdResponse = sendTxt2ImgRequest(params);
            log.info("Stable Diffusion API调用成功，返回图片数量: {}", sdResponse.getImages().size());
            
            // 4. 处理响应，上传图片到OSS
            Txt2ImgResponse result = processStableDiffusionResponse(sdResponse, params, request, startTime);
            
            // 5. 保存生成记录到数据库
            saveGenerationRecord(result, params, request, startTime);
            
            long totalTime = System.currentTimeMillis() - startTime;
            log.info("图片生成完成，任务ID: {}, 总耗时: {}ms", result.getTaskId(), totalTime);
            
            return result;
            
        } catch (Exception e) {
            log.error("图片生成失败", e);
            Txt2ImgResponse errorResponse = buildErrorResponse("图片生成失败: " + e.getMessage(), startTime);
            
            // 保存失败记录到数据库
            saveErrorRecord(errorResponse, request, startTime);
            
            return errorResponse;
        }
    }

    /**
     * 检查Stable Diffusion服务健康状态
     */
    @Override
    public boolean checkServiceHealth() {
        try {
            String url = apiBaseUrl + "/sdapi/v1/samplers";
            ResponseEntity<String> response = restTemplate.getForEntity(url, String.class);
            boolean isHealthy = response.getStatusCode() == HttpStatus.OK;
            log.info("Stable Diffusion服务健康检查: {}", isHealthy);
            return isHealthy;
        } catch (Exception e) {
            log.warn("Stable Diffusion服务健康检查失败: {}", e.getMessage());
            return false;
        }
    }

    /**
     * 发送Txt2Img请求到Stable Diffusion API
     */
    private StableDiffusionApiResponse sendTxt2ImgRequest(Text2ImgParameters params) {
        try {
            String url = apiBaseUrl + "/sdapi/v1/txt2img";
            
            log.info("发送Txt2Img请求到: {}", url);
            log.info("请求参数详情:");
            log.info("  prompt: {}", params.getPrompt());
            log.info("  negative_prompt: {}", params.getNegativePrompt());
            log.info("  steps: {}", params.getSteps());
            log.info("  sampler_name: {}", params.getSamplerName());
            log.info("  cfg_scale: {}", params.getCfgScale());
            log.info("  seed: {}", params.getSeed());
            log.info("  width: {}", params.getWidth());
            log.info("  height: {}", params.getHeight());
            log.info("  batch_size: {}", params.getBatchSize());
            log.info("  scheduler: {}", params.getScheduler());
            log.info("  denoising_strength: {}", params.getDenoisingStrength());
            log.info("  override_settings: {}", params.getOverrideSettings());
            
            // 设置请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            
            // 构建请求体
            Map<String, Object> requestBody = buildTxt2ImgRequestBody(params);
            
            // 记录完整的请求体
            log.info("完整请求体: {}", requestBody);
            
            // 创建请求实体
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            // 发送请求
            ResponseEntity<StableDiffusionApiResponse> response = restTemplate.exchange(
                url,
                HttpMethod.POST,
                requestEntity,
                StableDiffusionApiResponse.class
            );
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                log.info("Txt2Img请求成功，返回图片数量: {}", response.getBody().getImages().size());
                return response.getBody();
            } else {
                throw new RuntimeException("Txt2Img请求失败，状态码: " + response.getStatusCode());
            }
            
        } catch (Exception e) {
            log.error("发送Txt2Img请求失败", e);
            throw new RuntimeException("调用Stable Diffusion API失败: " + e.getMessage());
        }
    }

    /**
     * 构建Txt2Img请求体
     */
    private Map<String, Object> buildTxt2ImgRequestBody(Text2ImgParameters params) {
        try {
            log.debug("开始构建Txt2Img请求体，参数: {}", params);
            
            // 使用工具类构建请求体
            Map<String, Object> requestBody = text2imgConstructionUtil.buildRequestBody(params);
            
            log.debug("构建的请求体: {}", requestBody);
            return requestBody;
            
        } catch (Exception e) {
            log.error("构建请求体失败", e);
            throw new RuntimeException("构建请求体失败: " + e.getMessage());
        }
    }

    /**
     * 处理Stable Diffusion响应，上传图片到OSS
     */
    private Txt2ImgResponse processStableDiffusionResponse(StableDiffusionApiResponse sdResponse, 
                                                          Text2ImgParameters params,
                                                          Txt2ImgRequest request,
                                                          long startTime) {
        try {
            log.info("开始处理Stable Diffusion响应，图片数量: {}", sdResponse.getImages().size());
            
            List<String> imageUrls = new ArrayList<>();
            String taskId = UUID.randomUUID().toString();
            
            // 处理每张图片
            for (int i = 0; i < sdResponse.getImages().size(); i++) {
                String base64Image = sdResponse.getImages().get(i);
                
                // 生成文件名
                String fileName = generateImageFileName(taskId, i, params);
                
                // 上传到OSS
                String imageUrl = uploadBase64ImageToOSS(base64Image, fileName);
                imageUrls.add(imageUrl);
                
                log.info("图片{}上传成功", i);
            }
            
            // 解析生成信息
            Map<String, Object> infoMap = parseGenerationInfo(sdResponse.getInfo());
            
            // 构建响应
            Txt2ImgResponse result = new Txt2ImgResponse();
            result.setTaskId(taskId);
            result.setStatus("success");
            result.setImageUrls(imageUrls);
            result.setGenerationTime(System.currentTimeMillis() - startTime);
            result.setModelName(request.getModelName());
            result.setPrompt(params.getPrompt());
            result.setNegativePrompt(params.getNegativePrompt());
            result.setSeed(params.getSeed());
            result.setSamplerName(params.getSamplerName());
            result.setScheduler(params.getScheduler());
            result.setSteps(params.getSteps());
            result.setCfgScale(params.getCfgScale());
            result.setWidth(params.getWidth());
            result.setHeight(params.getHeight());
            result.setBatchSize(params.getBatchSize());
            result.setDenoisingStrength(params.getDenoisingStrength());
            
            // 从info中提取模型哈希
            if (infoMap != null && infoMap.containsKey("sd_model_hash")) {
                result.setModelHash((String) infoMap.get("sd_model_hash"));
            }
            
            log.info("图片处理完成，任务ID: {}, 生成时间: {}ms", taskId, result.getGenerationTime());
            return result;
            
        } catch (Exception e) {
            log.error("处理Stable Diffusion响应失败", e);
            throw new RuntimeException("图片处理失败: " + e.getMessage());
        }
    }

    /**
     * 保存成功生成记录到数据库
     */
    private void saveGenerationRecord(Txt2ImgResponse result, Text2ImgParameters params, 
                                    Txt2ImgRequest request, long startTime) {
        try {
            Text2Img record = new Text2Img();
            
            // 基本信息
            record.setTaskId(result.getTaskId());
            record.setPrompt(result.getPrompt());
            record.setNegativePrompt(result.getNegativePrompt());
            record.setSeed(result.getSeed());
            record.setSamplerName(result.getSamplerName());
            record.setSteps(result.getSteps());
            record.setCfgScale(result.getCfgScale());
            record.setWidth(result.getWidth());
            record.setHeight(result.getHeight());
            record.setDenoisingStrength(result.getDenoisingStrength());
            record.setModelName(result.getModelName());
            record.setModelHash(result.getModelHash());
            record.setBatchSize(result.getBatchSize());
            record.setScheduler(result.getScheduler());

            // 状态信息
            record.setStatus(result.getStatus());
            record.setGenerationTime(result.getGenerationTime());
            record.setErrorMessage(null); // 成功时错误信息为空
            
            // 图片URL列表（JSON格式）
            if (result.getImageUrls() != null && !result.getImageUrls().isEmpty()) {
                try {
                    String imageUrlsJson = objectMapper.writeValueAsString(result.getImageUrls());
                    record.setImageUrls(imageUrlsJson);
                } catch (Exception e) {
                    log.warn("序列化图片URL列表失败", e);
                    record.setImageUrls("[]");
                }
            } else {
                record.setImageUrls("[]");
            }
            
            // 默认值设置
            record.setRestoreFaces(false);
            record.setTiling(false);
            record.setEnableHr(false);
            record.setHrScale(null);
            record.setHrUpscaler(null);
            record.setHrSecondPassSteps(null);

            // 其他字段
            record.setCreateTime(LocalDateTime.now());
            record.setUpdateTime(LocalDateTime.now());
            record.setDeleted(0);

            // 保存到数据库
            boolean saveResult = text2ImgMapper.insert(record) > 0;
            if (saveResult) {
                log.info("Text2Img生成记录保存成功，任务ID: {}, 数据库ID: {}", result.getTaskId(), record.getId());
            } else {
                log.warn("Text2Img生成记录保存失败，任务ID: {}", result.getTaskId());
            }
            
        } catch (Exception e) {
            log.error("保存Text2Img生成记录失败，任务ID: {}", result.getTaskId(), e);
            // 不抛出异常，避免影响主流程
        }
    }

    /**
     * 保存失败记录到数据库
     */
    private void saveErrorRecord(Txt2ImgResponse errorResponse, Txt2ImgRequest request, long startTime) {
        try {
            Text2Img record = new Text2Img();
            
            // 基本信息
            record.setTaskId(errorResponse.getTaskId());
            record.setPrompt(request.getPrompt());
            record.setNegativePrompt(request.getNegativePrompt());
            record.setSeed(null); // 失败时种子为空
            record.setSamplerName(null);
            record.setSteps(null);
            record.setCfgScale(null);
            record.setWidth(request.getWidth());
            record.setHeight(request.getHeight());
            record.setDenoisingStrength(null);
            record.setModelName(request.getModelName());
            record.setModelHash(null);
            record.setBatchSize(null);
            record.setScheduler(null);
            
            // 状态信息
            record.setStatus(errorResponse.getStatus());
            record.setGenerationTime(errorResponse.getGenerationTime());
            record.setErrorMessage(errorResponse.getErrorMessage());
            
            // 图片URL列表为空
            record.setImageUrls("[]");
            
            // 默认值设置
            record.setRestoreFaces(false);
            record.setTiling(false);
            record.setEnableHr(false);
            record.setHrScale(null);
            record.setHrUpscaler(null);
            record.setHrSecondPassSteps(null);
            
            // 其他字段
            record.setCreateTime(LocalDateTime.now());
            record.setUpdateTime(LocalDateTime.now());
            record.setDeleted(0);
            
            // 保存到数据库
            boolean saveResult = text2ImgMapper.insert(record) > 0;
            if (saveResult) {
                log.info("Text2Img失败记录保存成功，任务ID: {}, 数据库ID: {}", errorResponse.getTaskId(), record.getId());
            } else {
                log.warn("Text2Img失败记录保存失败，任务ID: {}", errorResponse.getTaskId());
            }
            
        } catch (Exception e) {
            log.error("保存Text2Img失败记录失败，任务ID: {}", errorResponse.getTaskId(), e);
            // 不抛出异常，避免影响主流程
        }
    }

    /**
     * 生成图片文件名
     */
    private String generateImageFileName(String taskId, int index, Text2ImgParameters params) {
        String timestamp = String.valueOf(System.currentTimeMillis());
        String seed = String.valueOf(params.getSeed());
        
        // 格式: text2img_images/yyyy/MM/dd/taskId_index_seed_timestamp.png
        String datePath = java.time.LocalDate.now().toString().replace("-", "/");
        return String.format("text2img_images/%s/%s_%d_%s_%s.png", 
                           datePath, taskId, index, seed, timestamp);
    }

    /**
     * 上传base64图片到OSS
     */
    private String uploadBase64ImageToOSS(String base64Image, String fileName) {
        try {
            // 解码base64数据
            byte[] imageData = Base64.getDecoder().decode(base64Image);
            
            // 上传到OSS
            String imageUrl = aliOssUtil.upload(imageData, fileName);

            return imageUrl;
            
        } catch (Exception e) {
            log.error("上传图片到OSS失败: {}", fileName, e);
            throw new RuntimeException("图片上传失败: " + e.getMessage());
        }
    }

    /**
     * 解析生成信息
     */
    private Map<String, Object> parseGenerationInfo(String infoJson) {
        try {
            @SuppressWarnings("unchecked")
            Map<String, Object> infoMap = objectMapper.readValue(infoJson, Map.class);
            return infoMap;
        } catch (Exception e) {
            log.warn("解析生成信息失败: {}", infoJson, e);
            return null;
        }
    }

    /**
     * 构建错误响应
     */
    private Txt2ImgResponse buildErrorResponse(String errorMessage, long startTime) {
        Txt2ImgResponse response = new Txt2ImgResponse();
        response.setTaskId(UUID.randomUUID().toString());
        response.setStatus("failed");
        response.setErrorMessage(errorMessage);
        response.setGenerationTime(System.currentTimeMillis() - startTime);
        response.setImageUrls(new ArrayList<>());
        return response;
    }

    @Override
    public PageDTO<Text2Img> getPageRecords(Integer page, Integer size, String status, String modelName) {
        Page<Text2Img> pageParam = new Page<>(page, size);
        
        QueryWrapper<Text2Img> queryWrapper = new QueryWrapper<>();
        queryWrapper.orderByDesc("create_time");
        
        if (status != null && !status.trim().isEmpty()) {
            queryWrapper.eq("status", status);
        }
        
        if (modelName != null && !modelName.trim().isEmpty()) {
            queryWrapper.eq("model_name", modelName);
        }
        
        Page<Text2Img> result = text2ImgMapper.selectPage(pageParam, queryWrapper);
        return PageDTO.of(result, Text2Img.class);
    }
    
    @Override
    public Text2Img getByTaskId(String taskId) {
        QueryWrapper<Text2Img> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("task_id", taskId);
        return text2ImgMapper.selectOne(queryWrapper);
    }
    
    @Override
    public Map<String, Object> getStatistics() {
        // 总记录数
        Integer totalCount = text2ImgMapper.selectCount(null);
        
        // 成功记录数
        QueryWrapper<Text2Img> successWrapper = new QueryWrapper<>();
        successWrapper.eq("status", "success");
        Integer successCount = text2ImgMapper.selectCount(successWrapper);
        
        // 失败记录数
        QueryWrapper<Text2Img> failedWrapper = new QueryWrapper<>();
        failedWrapper.eq("status", "failed");
        Integer failedCount = text2ImgMapper.selectCount(failedWrapper);
        
        // 平均生成时间
        QueryWrapper<Text2Img> avgTimeWrapper = new QueryWrapper<>();
        avgTimeWrapper.eq("status", "success")
                      .select("AVG(generation_time) as avg_time");
        Map<String, Object> avgTimeResult = text2ImgMapper.selectMaps(avgTimeWrapper).isEmpty() ? null : text2ImgMapper.selectMaps(avgTimeWrapper).get(0);
        
        // 模型使用统计
        QueryWrapper<Text2Img> modelWrapper = new QueryWrapper<>();
        modelWrapper.select("model_name", "COUNT(*) as count")
                    .groupBy("model_name");
        List<Map<String, Object>> modelStats = text2ImgMapper.selectMaps(modelWrapper);
        
        // 构建统计结果
        Map<String, Object> statistics = new HashMap<>();
        statistics.put("totalCount", totalCount);
        statistics.put("successCount", successCount);
        statistics.put("failedCount", failedCount);
        statistics.put("successRate", totalCount > 0 ? (double) successCount / totalCount * 100 : 0);
        statistics.put("modelStats", modelStats);
        
        if (avgTimeResult != null) {
            statistics.put("avgGenerationTime", avgTimeResult.get("avg_time"));
        }
        
        return statistics;
    }
    
    @Override
    public boolean deleteByTaskId(String taskId) {
        QueryWrapper<Text2Img> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("task_id", taskId);
        int result = text2ImgMapper.delete(queryWrapper);
        return result > 0;
    }
} 